import os
import shutil
import sys
sys.path.append("./faceToolBox/lib")
from facedetect import facedetect
import cv2
import time
import numpy as np


def main():
    cnt=0

    fd_algo = "tf_ssd"
    videos_dir = "videos"
    result_dir = "valid_frame"

    if not os.path.exists(result_dir):
        os.makedirs(result_dir)

    face_fd = facedetect(fd_algo)

    for root, _, filenames in os.walk(videos_dir):
        for index, filename in enumerate(filenames) :
            videopath = os.path.join(root, filename)
            print(videopath)
            video = cv2.VideoCapture(videopath)

            save_dir = os.path.join(result_dir, filename)
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
                
            while True:
                success, frame = video.read()  

                if (success):
                    # cv2.imshow("test", frame)
                    # cv2.waitKey(1)
                    rgb_img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    rects, alignimgs= face_fd.findfaces(rgb_img, align=0)
                    if len(alignimgs)>0:
                        cv2.imwrite("%s/%d.jpg"%(save_dir, cnt), frame)
                        cnt+=1
                else:
                    break

    print("process down")
            


if __name__ =="__main__":
    main()